{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "### This note book is designed to determine the number of parameters of BiSeNetV2\n",
        "### Source of code: https://github.com/CoinCheung/BiSeNet/blob/master/lib/models/bisenetv2.py\n",
        "#BiSeNetV2 pytoruch version"
      ],
      "metadata": {
        "id": "22WpbQzLoj7I"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.utils.model_zoo as modelzoo\n",
        "\n",
        "backbone_url = 'https://github.com/CoinCheung/BiSeNet/releases/download/0.0.0/backbone_v2.pth'\n",
        "\n",
        "\n",
        "class ConvBNReLU(nn.Module):\n",
        "\n",
        "    def __init__(self, in_chan, out_chan, ks=3, stride=1, padding=1,\n",
        "                 dilation=1, groups=1, bias=False):\n",
        "        super(ConvBNReLU, self).__init__()\n",
        "        self.conv = nn.Conv2d(\n",
        "                in_chan, out_chan, kernel_size=ks, stride=stride,\n",
        "                padding=padding, dilation=dilation,\n",
        "                groups=groups, bias=bias)\n",
        "        self.bn = nn.BatchNorm2d(out_chan)\n",
        "        self.relu = nn.ReLU(inplace=True)\n",
        "\n",
        "    def forward(self, x):\n",
        "        feat = self.conv(x)\n",
        "        feat = self.bn(feat)\n",
        "        feat = self.relu(feat)\n",
        "        return feat\n",
        "\n",
        "\n",
        "class UpSample(nn.Module):\n",
        "\n",
        "    def __init__(self, n_chan, factor=2):\n",
        "        super(UpSample, self).__init__()\n",
        "        out_chan = n_chan * factor * factor\n",
        "        self.proj = nn.Conv2d(n_chan, out_chan, 1, 1, 0)\n",
        "        self.up = nn.PixelShuffle(factor)\n",
        "        self.init_weight()\n",
        "\n",
        "    def forward(self, x):\n",
        "        feat = self.proj(x)\n",
        "        feat = self.up(feat)\n",
        "        return feat\n",
        "\n",
        "    def init_weight(self):\n",
        "        nn.init.xavier_normal_(self.proj.weight, gain=1.)\n",
        "\n",
        "\n",
        "\n",
        "class DetailBranch(nn.Module):\n",
        "\n",
        "    def __init__(self):\n",
        "        super(DetailBranch, self).__init__()\n",
        "        self.S1 = nn.Sequential(\n",
        "            ConvBNReLU(3, 64, 3, stride=2),\n",
        "            ConvBNReLU(64, 64, 3, stride=1),\n",
        "        )\n",
        "        self.S2 = nn.Sequential(\n",
        "            ConvBNReLU(64, 64, 3, stride=2),\n",
        "            ConvBNReLU(64, 64, 3, stride=1),\n",
        "            ConvBNReLU(64, 64, 3, stride=1),\n",
        "        )\n",
        "        self.S3 = nn.Sequential(\n",
        "            ConvBNReLU(64, 128, 3, stride=2),\n",
        "            ConvBNReLU(128, 128, 3, stride=1),\n",
        "            ConvBNReLU(128, 128, 3, stride=1),\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        feat = self.S1(x)\n",
        "        feat = self.S2(feat)\n",
        "        feat = self.S3(feat)\n",
        "        return feat\n",
        "\n",
        "\n",
        "class StemBlock(nn.Module):\n",
        "\n",
        "    def __init__(self):\n",
        "        super(StemBlock, self).__init__()\n",
        "        self.conv = ConvBNReLU(3, 16, 3, stride=2)\n",
        "        self.left = nn.Sequential(\n",
        "            ConvBNReLU(16, 8, 1, stride=1, padding=0),\n",
        "            ConvBNReLU(8, 16, 3, stride=2),\n",
        "        )\n",
        "        self.right = nn.MaxPool2d(\n",
        "            kernel_size=3, stride=2, padding=1, ceil_mode=False)\n",
        "        self.fuse = ConvBNReLU(32, 16, 3, stride=1)\n",
        "\n",
        "    def forward(self, x):\n",
        "        feat = self.conv(x)\n",
        "        feat_left = self.left(feat)\n",
        "        feat_right = self.right(feat)\n",
        "        feat = torch.cat([feat_left, feat_right], dim=1)\n",
        "        feat = self.fuse(feat)\n",
        "        return feat\n",
        "\n",
        "\n",
        "class CEBlock(nn.Module):\n",
        "\n",
        "    def __init__(self):\n",
        "        super(CEBlock, self).__init__()\n",
        "        self.bn = nn.BatchNorm2d(128)\n",
        "        self.conv_gap = ConvBNReLU(128, 128, 1, stride=1, padding=0)\n",
        "        #TODO: in paper here is naive conv2d, no bn-relu\n",
        "        self.conv_last = ConvBNReLU(128, 128, 3, stride=1)\n",
        "\n",
        "    def forward(self, x):\n",
        "        feat = torch.mean(x, dim=(2, 3), keepdim=True)\n",
        "        feat = self.bn(feat)\n",
        "        feat = self.conv_gap(feat)\n",
        "        feat = feat + x\n",
        "        feat = self.conv_last(feat)\n",
        "        return feat\n",
        "\n",
        "\n",
        "class GELayerS1(nn.Module):\n",
        "\n",
        "    def __init__(self, in_chan, out_chan, exp_ratio=6):\n",
        "        super(GELayerS1, self).__init__()\n",
        "        mid_chan = in_chan * exp_ratio\n",
        "        self.conv1 = ConvBNReLU(in_chan, in_chan, 3, stride=1)\n",
        "        self.dwconv = nn.Sequential(\n",
        "            nn.Conv2d(\n",
        "                in_chan, mid_chan, kernel_size=3, stride=1,\n",
        "                padding=1, groups=in_chan, bias=False),\n",
        "            nn.BatchNorm2d(mid_chan),\n",
        "            nn.ReLU(inplace=True), # not shown in paper\n",
        "        )\n",
        "        self.conv2 = nn.Sequential(\n",
        "            nn.Conv2d(\n",
        "                mid_chan, out_chan, kernel_size=1, stride=1,\n",
        "                padding=0, bias=False),\n",
        "            nn.BatchNorm2d(out_chan),\n",
        "        )\n",
        "        self.conv2[1].last_bn = True\n",
        "        self.relu = nn.ReLU(inplace=True)\n",
        "\n",
        "    def forward(self, x):\n",
        "        feat = self.conv1(x)\n",
        "        feat = self.dwconv(feat)\n",
        "        feat = self.conv2(feat)\n",
        "        feat = feat + x\n",
        "        feat = self.relu(feat)\n",
        "        return feat\n",
        "\n",
        "\n",
        "class GELayerS2(nn.Module):\n",
        "\n",
        "    def __init__(self, in_chan, out_chan, exp_ratio=6):\n",
        "        super(GELayerS2, self).__init__()\n",
        "        mid_chan = in_chan * exp_ratio\n",
        "        self.conv1 = ConvBNReLU(in_chan, in_chan, 3, stride=1)\n",
        "        self.dwconv1 = nn.Sequential(\n",
        "            nn.Conv2d(\n",
        "                in_chan, mid_chan, kernel_size=3, stride=2,\n",
        "                padding=1, groups=in_chan, bias=False),\n",
        "            nn.BatchNorm2d(mid_chan),\n",
        "        )\n",
        "        self.dwconv2 = nn.Sequential(\n",
        "            nn.Conv2d(\n",
        "                mid_chan, mid_chan, kernel_size=3, stride=1,\n",
        "                padding=1, groups=mid_chan, bias=False),\n",
        "            nn.BatchNorm2d(mid_chan),\n",
        "            nn.ReLU(inplace=True), # not shown in paper\n",
        "        )\n",
        "        self.conv2 = nn.Sequential(\n",
        "            nn.Conv2d(\n",
        "                mid_chan, out_chan, kernel_size=1, stride=1,\n",
        "                padding=0, bias=False),\n",
        "            nn.BatchNorm2d(out_chan),\n",
        "        )\n",
        "        self.conv2[1].last_bn = True\n",
        "        self.shortcut = nn.Sequential(\n",
        "                nn.Conv2d(\n",
        "                    in_chan, in_chan, kernel_size=3, stride=2,\n",
        "                    padding=1, groups=in_chan, bias=False),\n",
        "                nn.BatchNorm2d(in_chan),\n",
        "                nn.Conv2d(\n",
        "                    in_chan, out_chan, kernel_size=1, stride=1,\n",
        "                    padding=0, bias=False),\n",
        "                nn.BatchNorm2d(out_chan),\n",
        "        )\n",
        "        self.relu = nn.ReLU(inplace=True)\n",
        "\n",
        "    def forward(self, x):\n",
        "        feat = self.conv1(x)\n",
        "        feat = self.dwconv1(feat)\n",
        "        feat = self.dwconv2(feat)\n",
        "        feat = self.conv2(feat)\n",
        "        shortcut = self.shortcut(x)\n",
        "        feat = feat + shortcut\n",
        "        feat = self.relu(feat)\n",
        "        return feat\n",
        "\n",
        "\n",
        "class SegmentBranch(nn.Module):\n",
        "\n",
        "    def __init__(self):\n",
        "        super(SegmentBranch, self).__init__()\n",
        "        self.S1S2 = StemBlock()\n",
        "        self.S3 = nn.Sequential(\n",
        "            GELayerS2(16, 32),\n",
        "            GELayerS1(32, 32),\n",
        "        )\n",
        "        self.S4 = nn.Sequential(\n",
        "            GELayerS2(32, 64),\n",
        "            GELayerS1(64, 64),\n",
        "        )\n",
        "        self.S5_4 = nn.Sequential(\n",
        "            GELayerS2(64, 128),\n",
        "            GELayerS1(128, 128),\n",
        "            GELayerS1(128, 128),\n",
        "            GELayerS1(128, 128),\n",
        "        )\n",
        "        self.S5_5 = CEBlock()\n",
        "\n",
        "    def forward(self, x):\n",
        "        feat2 = self.S1S2(x)\n",
        "        feat3 = self.S3(feat2)\n",
        "        feat4 = self.S4(feat3)\n",
        "        feat5_4 = self.S5_4(feat4)\n",
        "        feat5_5 = self.S5_5(feat5_4)\n",
        "        return feat2, feat3, feat4, feat5_4, feat5_5\n",
        "\n",
        "\n",
        "class BGALayer(nn.Module):\n",
        "\n",
        "    def __init__(self):\n",
        "        super(BGALayer, self).__init__()\n",
        "        self.left1 = nn.Sequential(\n",
        "            nn.Conv2d(\n",
        "                128, 128, kernel_size=3, stride=1,\n",
        "                padding=1, groups=128, bias=False),\n",
        "            nn.BatchNorm2d(128),\n",
        "            nn.Conv2d(\n",
        "                128, 128, kernel_size=1, stride=1,\n",
        "                padding=0, bias=False),\n",
        "        )\n",
        "        self.left2 = nn.Sequential(\n",
        "            nn.Conv2d(\n",
        "                128, 128, kernel_size=3, stride=2,\n",
        "                padding=1, bias=False),\n",
        "            nn.BatchNorm2d(128),\n",
        "            nn.AvgPool2d(kernel_size=3, stride=2, padding=1, ceil_mode=False)\n",
        "        )\n",
        "        self.right1 = nn.Sequential(\n",
        "            nn.Conv2d(\n",
        "                128, 128, kernel_size=3, stride=1,\n",
        "                padding=1, bias=False),\n",
        "            nn.BatchNorm2d(128),\n",
        "        )\n",
        "        self.right2 = nn.Sequential(\n",
        "            nn.Conv2d(\n",
        "                128, 128, kernel_size=3, stride=1,\n",
        "                padding=1, groups=128, bias=False),\n",
        "            nn.BatchNorm2d(128),\n",
        "            nn.Conv2d(\n",
        "                128, 128, kernel_size=1, stride=1,\n",
        "                padding=0, bias=False),\n",
        "        )\n",
        "        self.up1 = nn.Upsample(scale_factor=4)\n",
        "        self.up2 = nn.Upsample(scale_factor=4)\n",
        "        ##TODO: does this really has no relu?\n",
        "        self.conv = nn.Sequential(\n",
        "            nn.Conv2d(\n",
        "                128, 128, kernel_size=3, stride=1,\n",
        "                padding=1, bias=False),\n",
        "            nn.BatchNorm2d(128),\n",
        "            nn.ReLU(inplace=True), # not shown in paper\n",
        "        )\n",
        "\n",
        "    def forward(self, x_d, x_s):\n",
        "        dsize = x_d.size()[2:]\n",
        "        left1 = self.left1(x_d)\n",
        "        left2 = self.left2(x_d)\n",
        "        right1 = self.right1(x_s)\n",
        "        right2 = self.right2(x_s)\n",
        "        right1 = self.up1(right1)\n",
        "        left = left1 * torch.sigmoid(right1)\n",
        "        right = left2 * torch.sigmoid(right2)\n",
        "        right = self.up2(right)\n",
        "        out = self.conv(left + right)\n",
        "        return out\n",
        "\n",
        "\n",
        "\n",
        "class SegmentHead(nn.Module):\n",
        "\n",
        "    def __init__(self, in_chan, mid_chan, n_classes, up_factor=8, aux=True):\n",
        "        super(SegmentHead, self).__init__()\n",
        "        self.conv = ConvBNReLU(in_chan, mid_chan, 3, stride=1)\n",
        "        self.drop = nn.Dropout(0.1)\n",
        "        self.up_factor = up_factor\n",
        "\n",
        "        out_chan = n_classes\n",
        "        mid_chan2 = up_factor * up_factor if aux else mid_chan\n",
        "        up_factor = up_factor // 2 if aux else up_factor\n",
        "        self.conv_out = nn.Sequential(\n",
        "            nn.Sequential(\n",
        "                nn.Upsample(scale_factor=2),\n",
        "                ConvBNReLU(mid_chan, mid_chan2, 3, stride=1)\n",
        "                ) if aux else nn.Identity(),\n",
        "            nn.Conv2d(mid_chan2, out_chan, 1, 1, 0, bias=True),\n",
        "            nn.Upsample(scale_factor=up_factor, mode='bilinear', align_corners=False)\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        feat = self.conv(x)\n",
        "        feat = self.drop(feat)\n",
        "        feat = self.conv_out(feat)\n",
        "        return feat\n",
        "\n",
        "\n",
        "class BiSeNetV2(nn.Module):\n",
        "\n",
        "    def __init__(self, n_classes, aux_mode='train'):\n",
        "        super(BiSeNetV2, self).__init__()\n",
        "        self.aux_mode = aux_mode\n",
        "        self.detail = DetailBranch()\n",
        "        self.segment = SegmentBranch()\n",
        "        self.bga = BGALayer()\n",
        "\n",
        "        ## TODO: what is the number of mid chan ?\n",
        "        self.head = SegmentHead(128, 1024, n_classes, up_factor=8, aux=False)\n",
        "        if self.aux_mode == 'train':\n",
        "            self.aux2 = SegmentHead(16, 128, n_classes, up_factor=4)\n",
        "            self.aux3 = SegmentHead(32, 128, n_classes, up_factor=8)\n",
        "            self.aux4 = SegmentHead(64, 128, n_classes, up_factor=16)\n",
        "            self.aux5_4 = SegmentHead(128, 128, n_classes, up_factor=32)\n",
        "\n",
        "        self.init_weights()\n",
        "\n",
        "    def forward(self, x):\n",
        "        size = x.size()[2:]\n",
        "        feat_d = self.detail(x)\n",
        "        feat2, feat3, feat4, feat5_4, feat_s = self.segment(x)\n",
        "        feat_head = self.bga(feat_d, feat_s)\n",
        "\n",
        "        logits = self.head(feat_head)\n",
        "        if self.aux_mode == 'train':\n",
        "            logits_aux2 = self.aux2(feat2)\n",
        "            logits_aux3 = self.aux3(feat3)\n",
        "            logits_aux4 = self.aux4(feat4)\n",
        "            logits_aux5_4 = self.aux5_4(feat5_4)\n",
        "            return logits, logits_aux2, logits_aux3, logits_aux4, logits_aux5_4\n",
        "        elif self.aux_mode == 'eval':\n",
        "            return logits,\n",
        "        elif self.aux_mode == 'pred':\n",
        "            pred = logits.argmax(dim=1)\n",
        "            return pred\n",
        "        else:\n",
        "            raise NotImplementedError\n",
        "\n",
        "    def init_weights(self):\n",
        "        for name, module in self.named_modules():\n",
        "            if isinstance(module, (nn.Conv2d, nn.Linear)):\n",
        "                nn.init.kaiming_normal_(module.weight, mode='fan_out')\n",
        "                if not module.bias is None: nn.init.constant_(module.bias, 0)\n",
        "            elif isinstance(module, nn.modules.batchnorm._BatchNorm):\n",
        "                if hasattr(module, 'last_bn') and module.last_bn:\n",
        "                    nn.init.zeros_(module.weight)\n",
        "                else:\n",
        "                    nn.init.ones_(module.weight)\n",
        "                nn.init.zeros_(module.bias)\n",
        "        self.load_pretrain()\n",
        "\n",
        "\n",
        "    def load_pretrain(self):\n",
        "        state = modelzoo.load_url(backbone_url)\n",
        "        for name, child in self.named_children():\n",
        "            if name in state.keys():\n",
        "                child.load_state_dict(state[name], strict=True)\n",
        "\n",
        "    def get_params(self):\n",
        "        def add_param_to_list(mod, wd_params, nowd_params):\n",
        "            for param in mod.parameters():\n",
        "                if param.dim() == 1:\n",
        "                    nowd_params.append(param)\n",
        "                elif param.dim() == 4:\n",
        "                    wd_params.append(param)\n",
        "                else:\n",
        "                    print(name)\n",
        "\n",
        "        wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params = [], [], [], []\n",
        "        for name, child in self.named_children():\n",
        "            if 'head' in name or 'aux' in name:\n",
        "                add_param_to_list(child, lr_mul_wd_params, lr_mul_nowd_params)\n",
        "            else:\n",
        "                add_param_to_list(child, wd_params, nowd_params)\n",
        "        return wd_params, nowd_params, lr_mul_wd_params, lr_mul_nowd_params\n",
        "\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "\n",
        "    x = torch.randn(1, 3, 1024, 2048)\n",
        "    model = BiSeNetV2(n_classes=19)\n",
        "    #outs = model(x)\n",
        "    #for out in outs:\n",
        "    #    print(out.size())\n",
        "    #  print(logits.size())\n",
        "\n",
        "    for name, param in model.named_parameters():\n",
        "        if len(param.size()) == 1:\n",
        "            print(name)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Mr7JudFooebh",
        "outputId": "8aa46e5e-a37d-4630-cf74-efed835efb5d"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Downloading: \"https://github.com/CoinCheung/BiSeNet/releases/download/0.0.0/backbone_v2.pth\" to /root/.cache/torch/hub/checkpoints/backbone_v2.pth\n",
            "100%|██████████| 8.34M/8.34M [00:00<00:00, 283MB/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "detail.S1.0.bn.weight\n",
            "detail.S1.0.bn.bias\n",
            "detail.S1.1.bn.weight\n",
            "detail.S1.1.bn.bias\n",
            "detail.S2.0.bn.weight\n",
            "detail.S2.0.bn.bias\n",
            "detail.S2.1.bn.weight\n",
            "detail.S2.1.bn.bias\n",
            "detail.S2.2.bn.weight\n",
            "detail.S2.2.bn.bias\n",
            "detail.S3.0.bn.weight\n",
            "detail.S3.0.bn.bias\n",
            "detail.S3.1.bn.weight\n",
            "detail.S3.1.bn.bias\n",
            "detail.S3.2.bn.weight\n",
            "detail.S3.2.bn.bias\n",
            "segment.S1S2.conv.bn.weight\n",
            "segment.S1S2.conv.bn.bias\n",
            "segment.S1S2.left.0.bn.weight\n",
            "segment.S1S2.left.0.bn.bias\n",
            "segment.S1S2.left.1.bn.weight\n",
            "segment.S1S2.left.1.bn.bias\n",
            "segment.S1S2.fuse.bn.weight\n",
            "segment.S1S2.fuse.bn.bias\n",
            "segment.S3.0.conv1.bn.weight\n",
            "segment.S3.0.conv1.bn.bias\n",
            "segment.S3.0.dwconv1.1.weight\n",
            "segment.S3.0.dwconv1.1.bias\n",
            "segment.S3.0.dwconv2.1.weight\n",
            "segment.S3.0.dwconv2.1.bias\n",
            "segment.S3.0.conv2.1.weight\n",
            "segment.S3.0.conv2.1.bias\n",
            "segment.S3.0.shortcut.1.weight\n",
            "segment.S3.0.shortcut.1.bias\n",
            "segment.S3.0.shortcut.3.weight\n",
            "segment.S3.0.shortcut.3.bias\n",
            "segment.S3.1.conv1.bn.weight\n",
            "segment.S3.1.conv1.bn.bias\n",
            "segment.S3.1.dwconv.1.weight\n",
            "segment.S3.1.dwconv.1.bias\n",
            "segment.S3.1.conv2.1.weight\n",
            "segment.S3.1.conv2.1.bias\n",
            "segment.S4.0.conv1.bn.weight\n",
            "segment.S4.0.conv1.bn.bias\n",
            "segment.S4.0.dwconv1.1.weight\n",
            "segment.S4.0.dwconv1.1.bias\n",
            "segment.S4.0.dwconv2.1.weight\n",
            "segment.S4.0.dwconv2.1.bias\n",
            "segment.S4.0.conv2.1.weight\n",
            "segment.S4.0.conv2.1.bias\n",
            "segment.S4.0.shortcut.1.weight\n",
            "segment.S4.0.shortcut.1.bias\n",
            "segment.S4.0.shortcut.3.weight\n",
            "segment.S4.0.shortcut.3.bias\n",
            "segment.S4.1.conv1.bn.weight\n",
            "segment.S4.1.conv1.bn.bias\n",
            "segment.S4.1.dwconv.1.weight\n",
            "segment.S4.1.dwconv.1.bias\n",
            "segment.S4.1.conv2.1.weight\n",
            "segment.S4.1.conv2.1.bias\n",
            "segment.S5_4.0.conv1.bn.weight\n",
            "segment.S5_4.0.conv1.bn.bias\n",
            "segment.S5_4.0.dwconv1.1.weight\n",
            "segment.S5_4.0.dwconv1.1.bias\n",
            "segment.S5_4.0.dwconv2.1.weight\n",
            "segment.S5_4.0.dwconv2.1.bias\n",
            "segment.S5_4.0.conv2.1.weight\n",
            "segment.S5_4.0.conv2.1.bias\n",
            "segment.S5_4.0.shortcut.1.weight\n",
            "segment.S5_4.0.shortcut.1.bias\n",
            "segment.S5_4.0.shortcut.3.weight\n",
            "segment.S5_4.0.shortcut.3.bias\n",
            "segment.S5_4.1.conv1.bn.weight\n",
            "segment.S5_4.1.conv1.bn.bias\n",
            "segment.S5_4.1.dwconv.1.weight\n",
            "segment.S5_4.1.dwconv.1.bias\n",
            "segment.S5_4.1.conv2.1.weight\n",
            "segment.S5_4.1.conv2.1.bias\n",
            "segment.S5_4.2.conv1.bn.weight\n",
            "segment.S5_4.2.conv1.bn.bias\n",
            "segment.S5_4.2.dwconv.1.weight\n",
            "segment.S5_4.2.dwconv.1.bias\n",
            "segment.S5_4.2.conv2.1.weight\n",
            "segment.S5_4.2.conv2.1.bias\n",
            "segment.S5_4.3.conv1.bn.weight\n",
            "segment.S5_4.3.conv1.bn.bias\n",
            "segment.S5_4.3.dwconv.1.weight\n",
            "segment.S5_4.3.dwconv.1.bias\n",
            "segment.S5_4.3.conv2.1.weight\n",
            "segment.S5_4.3.conv2.1.bias\n",
            "segment.S5_5.bn.weight\n",
            "segment.S5_5.bn.bias\n",
            "segment.S5_5.conv_gap.bn.weight\n",
            "segment.S5_5.conv_gap.bn.bias\n",
            "segment.S5_5.conv_last.bn.weight\n",
            "segment.S5_5.conv_last.bn.bias\n",
            "bga.left1.1.weight\n",
            "bga.left1.1.bias\n",
            "bga.left2.1.weight\n",
            "bga.left2.1.bias\n",
            "bga.right1.1.weight\n",
            "bga.right1.1.bias\n",
            "bga.right2.1.weight\n",
            "bga.right2.1.bias\n",
            "bga.conv.1.weight\n",
            "bga.conv.1.bias\n",
            "head.conv.bn.weight\n",
            "head.conv.bn.bias\n",
            "head.conv_out.1.bias\n",
            "aux2.conv.bn.weight\n",
            "aux2.conv.bn.bias\n",
            "aux2.conv_out.0.1.bn.weight\n",
            "aux2.conv_out.0.1.bn.bias\n",
            "aux2.conv_out.1.bias\n",
            "aux3.conv.bn.weight\n",
            "aux3.conv.bn.bias\n",
            "aux3.conv_out.0.1.bn.weight\n",
            "aux3.conv_out.0.1.bn.bias\n",
            "aux3.conv_out.1.bias\n",
            "aux4.conv.bn.weight\n",
            "aux4.conv.bn.bias\n",
            "aux4.conv_out.0.1.bn.weight\n",
            "aux4.conv_out.0.1.bn.bias\n",
            "aux4.conv_out.1.bias\n",
            "aux5_4.conv.bn.weight\n",
            "aux5_4.conv.bn.bias\n",
            "aux5_4.conv_out.0.1.bn.weight\n",
            "aux5_4.conv_out.0.1.bn.bias\n",
            "aux5_4.conv_out.1.bias\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install prettytable"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Bd8SF68lBxxr",
        "outputId": "efb179d9-8b16-48d2-9df6-cda0431f86a5"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: prettytable in /usr/local/lib/python3.10/dist-packages (0.7.2)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from prettytable import PrettyTable\n",
        "def count_parameters(model):\n",
        "        table = PrettyTable([\"Modules\", \"Parameters\"])\n",
        "        total_params = 0\n",
        "        for name, parameter in model.named_parameters():\n",
        "            if not parameter.requires_grad: continue\n",
        "            param = parameter.numel()\n",
        "            table.add_row([name, param])\n",
        "            total_params+=param\n",
        "        print(table)\n",
        "        print(f\"Total Trainable Params: {total_params}\")\n",
        "        return total_params\n",
        "\n",
        "count_parameters(model)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Sa3JGm5fByHP",
        "outputId": "2fbb97e9-5e41-4da0-812e-b33a6da983f9"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "+------------------------------------+------------+\n",
            "|              Modules               | Parameters |\n",
            "+------------------------------------+------------+\n",
            "|      detail.S1.0.conv.weight       |    1728    |\n",
            "|       detail.S1.0.bn.weight        |     64     |\n",
            "|        detail.S1.0.bn.bias         |     64     |\n",
            "|      detail.S1.1.conv.weight       |   36864    |\n",
            "|       detail.S1.1.bn.weight        |     64     |\n",
            "|        detail.S1.1.bn.bias         |     64     |\n",
            "|      detail.S2.0.conv.weight       |   36864    |\n",
            "|       detail.S2.0.bn.weight        |     64     |\n",
            "|        detail.S2.0.bn.bias         |     64     |\n",
            "|      detail.S2.1.conv.weight       |   36864    |\n",
            "|       detail.S2.1.bn.weight        |     64     |\n",
            "|        detail.S2.1.bn.bias         |     64     |\n",
            "|      detail.S2.2.conv.weight       |   36864    |\n",
            "|       detail.S2.2.bn.weight        |     64     |\n",
            "|        detail.S2.2.bn.bias         |     64     |\n",
            "|      detail.S3.0.conv.weight       |   73728    |\n",
            "|       detail.S3.0.bn.weight        |    128     |\n",
            "|        detail.S3.0.bn.bias         |    128     |\n",
            "|      detail.S3.1.conv.weight       |   147456   |\n",
            "|       detail.S3.1.bn.weight        |    128     |\n",
            "|        detail.S3.1.bn.bias         |    128     |\n",
            "|      detail.S3.2.conv.weight       |   147456   |\n",
            "|       detail.S3.2.bn.weight        |    128     |\n",
            "|        detail.S3.2.bn.bias         |    128     |\n",
            "|   segment.S1S2.conv.conv.weight    |    432     |\n",
            "|    segment.S1S2.conv.bn.weight     |     16     |\n",
            "|     segment.S1S2.conv.bn.bias      |     16     |\n",
            "|  segment.S1S2.left.0.conv.weight   |    128     |\n",
            "|   segment.S1S2.left.0.bn.weight    |     8      |\n",
            "|    segment.S1S2.left.0.bn.bias     |     8      |\n",
            "|  segment.S1S2.left.1.conv.weight   |    1152    |\n",
            "|   segment.S1S2.left.1.bn.weight    |     16     |\n",
            "|    segment.S1S2.left.1.bn.bias     |     16     |\n",
            "|   segment.S1S2.fuse.conv.weight    |    4608    |\n",
            "|    segment.S1S2.fuse.bn.weight     |     16     |\n",
            "|     segment.S1S2.fuse.bn.bias      |     16     |\n",
            "|   segment.S3.0.conv1.conv.weight   |    2304    |\n",
            "|    segment.S3.0.conv1.bn.weight    |     16     |\n",
            "|     segment.S3.0.conv1.bn.bias     |     16     |\n",
            "|   segment.S3.0.dwconv1.0.weight    |    864     |\n",
            "|   segment.S3.0.dwconv1.1.weight    |     96     |\n",
            "|    segment.S3.0.dwconv1.1.bias     |     96     |\n",
            "|   segment.S3.0.dwconv2.0.weight    |    864     |\n",
            "|   segment.S3.0.dwconv2.1.weight    |     96     |\n",
            "|    segment.S3.0.dwconv2.1.bias     |     96     |\n",
            "|    segment.S3.0.conv2.0.weight     |    3072    |\n",
            "|    segment.S3.0.conv2.1.weight     |     32     |\n",
            "|     segment.S3.0.conv2.1.bias      |     32     |\n",
            "|   segment.S3.0.shortcut.0.weight   |    144     |\n",
            "|   segment.S3.0.shortcut.1.weight   |     16     |\n",
            "|    segment.S3.0.shortcut.1.bias    |     16     |\n",
            "|   segment.S3.0.shortcut.2.weight   |    512     |\n",
            "|   segment.S3.0.shortcut.3.weight   |     32     |\n",
            "|    segment.S3.0.shortcut.3.bias    |     32     |\n",
            "|   segment.S3.1.conv1.conv.weight   |    9216    |\n",
            "|    segment.S3.1.conv1.bn.weight    |     32     |\n",
            "|     segment.S3.1.conv1.bn.bias     |     32     |\n",
            "|    segment.S3.1.dwconv.0.weight    |    1728    |\n",
            "|    segment.S3.1.dwconv.1.weight    |    192     |\n",
            "|     segment.S3.1.dwconv.1.bias     |    192     |\n",
            "|    segment.S3.1.conv2.0.weight     |    6144    |\n",
            "|    segment.S3.1.conv2.1.weight     |     32     |\n",
            "|     segment.S3.1.conv2.1.bias      |     32     |\n",
            "|   segment.S4.0.conv1.conv.weight   |    9216    |\n",
            "|    segment.S4.0.conv1.bn.weight    |     32     |\n",
            "|     segment.S4.0.conv1.bn.bias     |     32     |\n",
            "|   segment.S4.0.dwconv1.0.weight    |    1728    |\n",
            "|   segment.S4.0.dwconv1.1.weight    |    192     |\n",
            "|    segment.S4.0.dwconv1.1.bias     |    192     |\n",
            "|   segment.S4.0.dwconv2.0.weight    |    1728    |\n",
            "|   segment.S4.0.dwconv2.1.weight    |    192     |\n",
            "|    segment.S4.0.dwconv2.1.bias     |    192     |\n",
            "|    segment.S4.0.conv2.0.weight     |   12288    |\n",
            "|    segment.S4.0.conv2.1.weight     |     64     |\n",
            "|     segment.S4.0.conv2.1.bias      |     64     |\n",
            "|   segment.S4.0.shortcut.0.weight   |    288     |\n",
            "|   segment.S4.0.shortcut.1.weight   |     32     |\n",
            "|    segment.S4.0.shortcut.1.bias    |     32     |\n",
            "|   segment.S4.0.shortcut.2.weight   |    2048    |\n",
            "|   segment.S4.0.shortcut.3.weight   |     64     |\n",
            "|    segment.S4.0.shortcut.3.bias    |     64     |\n",
            "|   segment.S4.1.conv1.conv.weight   |   36864    |\n",
            "|    segment.S4.1.conv1.bn.weight    |     64     |\n",
            "|     segment.S4.1.conv1.bn.bias     |     64     |\n",
            "|    segment.S4.1.dwconv.0.weight    |    3456    |\n",
            "|    segment.S4.1.dwconv.1.weight    |    384     |\n",
            "|     segment.S4.1.dwconv.1.bias     |    384     |\n",
            "|    segment.S4.1.conv2.0.weight     |   24576    |\n",
            "|    segment.S4.1.conv2.1.weight     |     64     |\n",
            "|     segment.S4.1.conv2.1.bias      |     64     |\n",
            "|  segment.S5_4.0.conv1.conv.weight  |   36864    |\n",
            "|   segment.S5_4.0.conv1.bn.weight   |     64     |\n",
            "|    segment.S5_4.0.conv1.bn.bias    |     64     |\n",
            "|  segment.S5_4.0.dwconv1.0.weight   |    3456    |\n",
            "|  segment.S5_4.0.dwconv1.1.weight   |    384     |\n",
            "|   segment.S5_4.0.dwconv1.1.bias    |    384     |\n",
            "|  segment.S5_4.0.dwconv2.0.weight   |    3456    |\n",
            "|  segment.S5_4.0.dwconv2.1.weight   |    384     |\n",
            "|   segment.S5_4.0.dwconv2.1.bias    |    384     |\n",
            "|   segment.S5_4.0.conv2.0.weight    |   49152    |\n",
            "|   segment.S5_4.0.conv2.1.weight    |    128     |\n",
            "|    segment.S5_4.0.conv2.1.bias     |    128     |\n",
            "|  segment.S5_4.0.shortcut.0.weight  |    576     |\n",
            "|  segment.S5_4.0.shortcut.1.weight  |     64     |\n",
            "|   segment.S5_4.0.shortcut.1.bias   |     64     |\n",
            "|  segment.S5_4.0.shortcut.2.weight  |    8192    |\n",
            "|  segment.S5_4.0.shortcut.3.weight  |    128     |\n",
            "|   segment.S5_4.0.shortcut.3.bias   |    128     |\n",
            "|  segment.S5_4.1.conv1.conv.weight  |   147456   |\n",
            "|   segment.S5_4.1.conv1.bn.weight   |    128     |\n",
            "|    segment.S5_4.1.conv1.bn.bias    |    128     |\n",
            "|   segment.S5_4.1.dwconv.0.weight   |    6912    |\n",
            "|   segment.S5_4.1.dwconv.1.weight   |    768     |\n",
            "|    segment.S5_4.1.dwconv.1.bias    |    768     |\n",
            "|   segment.S5_4.1.conv2.0.weight    |   98304    |\n",
            "|   segment.S5_4.1.conv2.1.weight    |    128     |\n",
            "|    segment.S5_4.1.conv2.1.bias     |    128     |\n",
            "|  segment.S5_4.2.conv1.conv.weight  |   147456   |\n",
            "|   segment.S5_4.2.conv1.bn.weight   |    128     |\n",
            "|    segment.S5_4.2.conv1.bn.bias    |    128     |\n",
            "|   segment.S5_4.2.dwconv.0.weight   |    6912    |\n",
            "|   segment.S5_4.2.dwconv.1.weight   |    768     |\n",
            "|    segment.S5_4.2.dwconv.1.bias    |    768     |\n",
            "|   segment.S5_4.2.conv2.0.weight    |   98304    |\n",
            "|   segment.S5_4.2.conv2.1.weight    |    128     |\n",
            "|    segment.S5_4.2.conv2.1.bias     |    128     |\n",
            "|  segment.S5_4.3.conv1.conv.weight  |   147456   |\n",
            "|   segment.S5_4.3.conv1.bn.weight   |    128     |\n",
            "|    segment.S5_4.3.conv1.bn.bias    |    128     |\n",
            "|   segment.S5_4.3.dwconv.0.weight   |    6912    |\n",
            "|   segment.S5_4.3.dwconv.1.weight   |    768     |\n",
            "|    segment.S5_4.3.dwconv.1.bias    |    768     |\n",
            "|   segment.S5_4.3.conv2.0.weight    |   98304    |\n",
            "|   segment.S5_4.3.conv2.1.weight    |    128     |\n",
            "|    segment.S5_4.3.conv2.1.bias     |    128     |\n",
            "|       segment.S5_5.bn.weight       |    128     |\n",
            "|        segment.S5_5.bn.bias        |    128     |\n",
            "| segment.S5_5.conv_gap.conv.weight  |   16384    |\n",
            "|  segment.S5_5.conv_gap.bn.weight   |    128     |\n",
            "|   segment.S5_5.conv_gap.bn.bias    |    128     |\n",
            "| segment.S5_5.conv_last.conv.weight |   147456   |\n",
            "|  segment.S5_5.conv_last.bn.weight  |    128     |\n",
            "|   segment.S5_5.conv_last.bn.bias   |    128     |\n",
            "|         bga.left1.0.weight         |    1152    |\n",
            "|         bga.left1.1.weight         |    128     |\n",
            "|          bga.left1.1.bias          |    128     |\n",
            "|         bga.left1.2.weight         |   16384    |\n",
            "|         bga.left2.0.weight         |   147456   |\n",
            "|         bga.left2.1.weight         |    128     |\n",
            "|          bga.left2.1.bias          |    128     |\n",
            "|        bga.right1.0.weight         |   147456   |\n",
            "|        bga.right1.1.weight         |    128     |\n",
            "|         bga.right1.1.bias          |    128     |\n",
            "|        bga.right2.0.weight         |    1152    |\n",
            "|        bga.right2.1.weight         |    128     |\n",
            "|         bga.right2.1.bias          |    128     |\n",
            "|        bga.right2.2.weight         |   16384    |\n",
            "|         bga.conv.0.weight          |   147456   |\n",
            "|         bga.conv.1.weight          |    128     |\n",
            "|          bga.conv.1.bias           |    128     |\n",
            "|       head.conv.conv.weight        |  1179648   |\n",
            "|        head.conv.bn.weight         |    1024    |\n",
            "|         head.conv.bn.bias          |    1024    |\n",
            "|       head.conv_out.1.weight       |   19456    |\n",
            "|        head.conv_out.1.bias        |     19     |\n",
            "|       aux2.conv.conv.weight        |   18432    |\n",
            "|        aux2.conv.bn.weight         |    128     |\n",
            "|         aux2.conv.bn.bias          |    128     |\n",
            "|   aux2.conv_out.0.1.conv.weight    |   18432    |\n",
            "|    aux2.conv_out.0.1.bn.weight     |     16     |\n",
            "|     aux2.conv_out.0.1.bn.bias      |     16     |\n",
            "|       aux2.conv_out.1.weight       |    304     |\n",
            "|        aux2.conv_out.1.bias        |     19     |\n",
            "|       aux3.conv.conv.weight        |   36864    |\n",
            "|        aux3.conv.bn.weight         |    128     |\n",
            "|         aux3.conv.bn.bias          |    128     |\n",
            "|   aux3.conv_out.0.1.conv.weight    |   73728    |\n",
            "|    aux3.conv_out.0.1.bn.weight     |     64     |\n",
            "|     aux3.conv_out.0.1.bn.bias      |     64     |\n",
            "|       aux3.conv_out.1.weight       |    1216    |\n",
            "|        aux3.conv_out.1.bias        |     19     |\n",
            "|       aux4.conv.conv.weight        |   73728    |\n",
            "|        aux4.conv.bn.weight         |    128     |\n",
            "|         aux4.conv.bn.bias          |    128     |\n",
            "|   aux4.conv_out.0.1.conv.weight    |   294912   |\n",
            "|    aux4.conv_out.0.1.bn.weight     |    256     |\n",
            "|     aux4.conv_out.0.1.bn.bias      |    256     |\n",
            "|       aux4.conv_out.1.weight       |    4864    |\n",
            "|        aux4.conv_out.1.bias        |     19     |\n",
            "|      aux5_4.conv.conv.weight       |   147456   |\n",
            "|       aux5_4.conv.bn.weight        |    128     |\n",
            "|        aux5_4.conv.bn.bias         |    128     |\n",
            "|  aux5_4.conv_out.0.1.conv.weight   |  1179648   |\n",
            "|   aux5_4.conv_out.0.1.bn.weight    |    1024    |\n",
            "|    aux5_4.conv_out.0.1.bn.bias     |    1024    |\n",
            "|      aux5_4.conv_out.1.weight      |   19456    |\n",
            "|       aux5_4.conv_out.1.bias       |     19     |\n",
            "+------------------------------------+------------+\n",
            "Total Trainable Params: 5231487\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "5231487"
            ]
          },
          "metadata": {},
          "execution_count": 3
        }
      ]
    }
  ]
}